{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d7fd6e26",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "# %load_ext nb_black"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# CUDA_VISIBLE_DEVICES=0\n",
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4f905a29",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[32m2024-02-02 14:22:33.271\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mmain_dense\u001b[0m:\u001b[36mload_models\u001b[0m:\u001b[36m835\u001b[0m - \u001b[1mLoading biencoder model\u001b[0m\n",
      "\u001b[32m2024-02-02 14:22:49.971\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mmain_dense\u001b[0m:\u001b[36mload_models\u001b[0m:\u001b[36m873\u001b[0m - \u001b[1mLoading candidate entities\u001b[0m\n",
      "parsing id2 title/text/wikidata: 100%|██████████| 5903527/5903527 [00:29<00:00, 200664.13it/s]\n"
     ]
    }
   ],
   "source": [
    "import main_dense as main_dense\n",
    "import argparse\n",
    "from loguru import logger\n",
    "\n",
    "models_path = \"ELQ_data/\" # the path where you stored the ELQ models\n",
    "\n",
    "config = {\n",
    "    \"interactive\": False,\n",
    "    \"biencoder_model\": models_path+\"elq_wiki_large.bin\",\n",
    "    \"biencoder_config\": models_path+\"elq_large_params.txt\",\n",
    "    \"cand_token_ids_path\": models_path+\"entity_token_ids_128.t7\",\n",
    "    \"entity_catalogue\": models_path+\"entity.jsonl\",\n",
    "    \"entity_encoding\": models_path+\"all_entities_large.t7\",\n",
    "    \"output_path\": \"logs/\", # logging directory\n",
    "    \"faiss_index\": \"hnsw\",\n",
    "    \"index_path\": models_path+\"faiss_hnsw_index.pkl\",\n",
    "    \"num_cand_mentions\": 10,\n",
    "    \"num_cand_entities\": 10,\n",
    "    \"threshold_type\": \"joint\",\n",
    "    \"threshold\": -4.5,\n",
    "    \"use_cuda\":True,\n",
    "}\n",
    "\n",
    "args = argparse.Namespace(**config)\n",
    "\n",
    "models = main_dense.load_models(args, logger=logger)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "eab5758b",
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:00<00:00,  7.94it/s]\n"
     ]
    }
   ],
   "source": [
    "data_to_link = [{\n",
    "                    \"id\": \"qwe\",\n",
    "                    \"text\": \"paris is capital of which country?\".lower(),\n",
    "                },\n",
    "                {\n",
    "                    \"id\": 1,\n",
    "                    \"text\": \"paris is great granddaughter of whom?\".lower(),\n",
    "                },\n",
    "                {\n",
    "                    \"id\": 2,\n",
    "                    \"text\": \"who discovered o in the periodic table?\".lower(),\n",
    "                },\n",
    "                ]\n",
    "args.threshold = -1.5\n",
    "args.test_mentions= \"./\"\n",
    "predictions = main_dense.run(args, None, *models, test_data=data_to_link)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a95cce6c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'id': 'qwe',\n",
       "  'text': 'paris is capital of which country?',\n",
       "  'scores': [-0.9338265061378479, -3.9720418453216553],\n",
       "  'pred_tuples_string': [['Paris', 'paris'], ['Capital city', 'capital']],\n",
       "  'pred_triples': [('11245', 0, 1), ('100454', 2, 3)],\n",
       "  'tokens': [3000, 2003, 3007, 1997, 2029, 2406, 1029]},\n",
       " {'id': 1,\n",
       "  'text': 'paris is great granddaughter of whom?',\n",
       "  'scores': [-3.798147439956665],\n",
       "  'pred_tuples_string': [['Paris Hilton', 'paris']],\n",
       "  'pred_triples': [('1610293', 0, 1)],\n",
       "  'tokens': [3000, 2003, 2307, 12787, 1997, 3183, 1029]},\n",
       " {'id': 2,\n",
       "  'text': 'who discovered o in the periodic table?',\n",
       "  'scores': [-0.5392034649848938, -3.703461170196533],\n",
       "  'pred_tuples_string': [['Periodic table', 'periodic table'],\n",
       "   ['Oxygen', 'o']],\n",
       "  'pred_triples': [('11282', 5, 7), ('10935', 2, 3)],\n",
       "  'tokens': [2040, 3603, 1051, 1999, 1996, 15861, 2795, 1029]}]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "predictions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## small data webqsp and cwq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/xionggm/codes/LLM_KGQA/save-anno:\n",
      "cwq  kqapro  metaqa  webqsp\n",
      "\n",
      "/home/xionggm/codes/LLM_KGQA/save-anno-clean:\n",
      "cwq  kqapro  metaqa  webqsp\n"
     ]
    }
   ],
   "source": [
    "!ls /home/xionggm/codes/LLM_KGQA/save-anno-clean/webqsp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "from glob import glob\n",
    "import json\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "900"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = []\n",
    "paths = glob(\"/home/xionggm/codes/LLM_KGQA/data/webqsp/test/*.json\") + glob(\"/home/xionggm/codes/LLM_KGQA/data/cwq/test/*.json\")\n",
    "for p in paths:\n",
    "    d = json.load(open(p))\n",
    "    data.extend(d)\n",
    "len(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1200"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "paths = glob(\"/home/xionggm/codes/LLM_KGQA/save-anno-clean/webqsp/**/*.json\") + glob(\"/home/xionggm/codes/LLM_KGQA/save-anno-clean/cwq/**/*.json\")\n",
    "for p in paths:\n",
    "    d = json.load(open(p))\n",
    "    data.append(d)\n",
    "len(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1200"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "all_ques = [{\"id\": d[\"id\"], \"text\": d[\"question\"].lower()} for d in data]\n",
    "len(all_ques)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'id': 'WebQTest-516.P0', 'text': 'what is the government system of malaysia?'}"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "all_ques[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[32m2024-02-02 14:57:16.432\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mmain_dense\u001b[0m:\u001b[36mrun\u001b[0m:\u001b[36m1025\u001b[0m - \u001b[1mPreparing data for biencoder\u001b[0m\n",
      "\u001b[32m2024-02-02 14:57:16.639\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mmain_dense\u001b[0m:\u001b[36mrun\u001b[0m:\u001b[36m1046\u001b[0m - \u001b[1mRunning biencoder...\u001b[0m\n",
      "100%|██████████| 150/150 [00:37<00:00,  3.97it/s]\n",
      "\u001b[32m2024-02-02 14:57:54.427\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mmain_dense\u001b[0m:\u001b[36mrun\u001b[0m:\u001b[36m1069\u001b[0m - \u001b[1mFinished running biencoder\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "args.threshold = -4.5\n",
    "args.test_mentions= \"./\"\n",
    "predictions = main_dense.run(args, logger, *models, test_data=all_ques)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "for p in predictions:\n",
    "    p.pop(\"tokens\", None)\n",
    "    p.pop(\"pred_triples\", None)\n",
    "json.dump(predictions, open(\"elq-webqsp-cwq-1200-threshold-neg4.5.json\", \"w\"),indent=4, ensure_ascii=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b858c6e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from common_utils import read_jsonl,read_json,save_to_json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "1470f7a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "ents = read_jsonl(\"blink_elq_data/elq/models/entity.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "49a4ac46",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'text': \" Paris () is the capital and most populous city of France, with an area of and an official estimated population of 2,140,526 residents as of 1 January 2019. Since the 17th century, Paris has been one of Europe's major centres of finance, diplomacy, commerce, fashion, science, as well as the arts. The City of Paris is the centre and seat of government of the Ile-de-France, or Paris Region, which has an estimated official 2019 population of 12,213,364, or about 18 percent of the population of France. The Paris Region had a GDP of €709 billion ($808 billion) in 2017. According to the Economist Intelligence Unit Worldwide Cost of Living Survey in 2018, Paris was the second most expensive city in the world, after Singapore, and ahead of Zurich, Hong Kong, Oslo and Geneva. Another source ranked Paris as most expensive, on a par with Singapore and Hong Kong, in 2018. The city is a major railway, highway, and air-transport hub served by two international airports: Paris-Charles de Gaulle (the second busiest airport in Europe) and Paris-Orly. Opened in 1900, the city's subway system, the Paris Metro, serves 5.23 million passengers daily, and is the second busiest metro system in Europe after Moscow Metro. Gare du Nord is the 24th busiest railway station in the world, but the first located outside Japan, with 262 million passengers in 2015.  Paris is especially known for its museums and architectural landmarks: the Louvre was the most visited art museum in the world in 2018, with 10.2 million visitors.\",\n",
       " 'idx': 'https://en.wikipedia.org/wiki?curid=22989',\n",
       " 'title': 'Paris',\n",
       " 'entity': 'Paris',\n",
       " 'kb_idx': 'Q90'}"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ents[11245]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "51f385d1",
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "def el(datas):\n",
    "    args.threshold = -4.5\n",
    "    args.test_mentions= \"./\"\n",
    "    predictions = main_dense.run(args, None, *models, test_data=datas)\n",
    "    return predictions\n",
    "def add_qid(predictions):\n",
    "    for item in predictions:\n",
    "        _wikis = []\n",
    "        for p in item[\"pred_triples\"]:\n",
    "            _w = ents[int(p[0])].get(\"kb_idx\",None)\n",
    "            _wikis.append(_w)\n",
    "        item[\"wikiids\"] = _wikis\n",
    "    return predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "035885e8",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "b5911d1c",
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "datas = [{\n",
    "                    \"id\": \"qwe\",\n",
    "                    \"text\": \"paris is capital of which country?\".lower(),\n",
    "                },\n",
    "                {\n",
    "                    \"id\": 1,\n",
    "                    \"text\": \"paris is great granddaughter of whom?\".lower(),\n",
    "                },\n",
    "                {\n",
    "                    \"id\": 2,\n",
    "                    \"text\": \"who discovered o in the periodic table?\".lower(),\n",
    "                },\n",
    "                ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "585cc308",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "no cuda:  False\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:00<00:00, 10.66it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'id': 'qwe',\n",
       "  'text': 'paris is capital of which country?',\n",
       "  'scores': [-0.9336556196212769, -3.971074342727661],\n",
       "  'pred_tuples_string': [['Paris', 'paris'], ['Capital city', 'capital']],\n",
       "  'pred_triples': [('11245', 0, 1), ('100454', 2, 3)],\n",
       "  'tokens': [3000, 2003, 3007, 1997, 2029, 2406, 1029]},\n",
       " {'id': 1,\n",
       "  'text': 'paris is great granddaughter of whom?',\n",
       "  'scores': [-3.8006019592285156],\n",
       "  'pred_tuples_string': [['Paris Hilton', 'paris']],\n",
       "  'pred_triples': [('1610293', 0, 1)],\n",
       "  'tokens': [3000, 2003, 2307, 12787, 1997, 3183, 1029]},\n",
       " {'id': 2,\n",
       "  'text': 'who discovered o in the periodic table?',\n",
       "  'scores': [-0.5388863682746887, -3.7036046981811523],\n",
       "  'pred_tuples_string': [['Periodic table', 'periodic table'],\n",
       "   ['Oxygen', 'o']],\n",
       "  'pred_triples': [('11282', 5, 7), ('10935', 2, 3)],\n",
       "  'tokens': [2040, 3603, 1051, 1999, 1996, 15861, 2795, 1029]}]"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "el(datas)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d625757e",
   "metadata": {},
   "source": [
    "## grailqa"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "b01d60d8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "64331"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "grailqas = []\n",
    "ids = {}\n",
    "for p in [\"/home/jimx/codes/rng-kbqa/GrailQA/outputs/grailqa_v1.0_train.json\",\n",
    "         \"/home/jimx/codes/rng-kbqa/GrailQA/outputs/grailqa_v1.0_dev.json\",\n",
    "         \"/home/jimx/codes/rng-kbqa/GrailQA/outputs/grailqa_v1.0_test.json\"]:\n",
    "    for item in read_json(p):\n",
    "        _id = item[\"qid\"]\n",
    "        assert _id not in ids\n",
    "        ids[_id] = 1\n",
    "        tmp = {\"id\":_id,\"text\":item[\"question\"].lower()}\n",
    "        grailqas.append(tmp)\n",
    "len(grailqas)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1249294c",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "no cuda:  False\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 48%|████▊     | 3835/8042 [17:43<18:52,  3.71it/s]"
     ]
    }
   ],
   "source": [
    "preds = el(grailqas[:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d014fa1",
   "metadata": {},
   "outputs": [],
   "source": [
    "preds=add_qid(preds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e3698fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "save_to_json(preds,\"grailqa_ents.json\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e227b46",
   "metadata": {},
   "source": [
    "## wqcwq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c59f4e81",
   "metadata": {},
   "outputs": [],
   "source": [
    "wqcwqs = []\n",
    "ids = {}\n",
    "for p in [\n",
    "    \"/home/jimx/f/datasets/complexwebquestions_V1_1/ComplexWebQuestions_train.json\",\n",
    "    \"/home/jimx/f/datasets/complexwebquestions_V1_1/ComplexWebQuestions_dev.json\",\n",
    "    \"/home/jimx/f/datasets/complexwebquestions_V1_1/ComplexWebQuestions_test.json\",\n",
    "    ]:\n",
    "    for item in read_json(p):\n",
    "        _id = item[\"ID\"]\n",
    "        assert _id not in ids\n",
    "        ids[_id] = 1\n",
    "        tmp = {\"id\":_id,\"text\":item[\"question\"].lower()}\n",
    "        wqcwqs.append(tmp)\n",
    "len(wqcwqs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "067c1e40",
   "metadata": {},
   "outputs": [],
   "source": [
    "wqcwq_preds = el(wqcwqs[:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e13839a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(wqcwq_preds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62d63e1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "wqcwq_preds=add_qid(wqcwq_preds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44530145",
   "metadata": {},
   "outputs": [],
   "source": [
    "save_to_json(wqcwq_preds,\"wqcwq_ents.json\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aee9b09e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "451554e0",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {
    "height": "calc(100% - 180px)",
    "left": "10px",
    "top": "150px",
    "width": "298.466px"
   },
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
